from langchain.text_splitter import CharacterTextSplitter
from langchain.prompts import ChatPromptTemplate
from langchain.schema import StrOutputParser
from langchain.schema.runnable import RunnablePassthrough
from typing import Optional
import logging
from .custom_chat_model import DeepSeekChatModel

logger = logging.getLogger(__name__)

class LangChainService:
    def __init__(self):
        self.chat_model = DeepSeekChatModel()
        self.text_splitter = CharacterTextSplitter(
            separator="\n",
            chunk_size=4000,
            chunk_overlap=200,
            length_function=len
        )

    def process_text(self, text: str, prompt_type: str, 
                    target_language: Optional[str] = None,
                    summary_length: Optional[int] = None) -> str:
        """处理文本"""
        try:
            logger.info(f"Processing text with type: {prompt_type}")
            # 将文本分割成块
            docs = self.text_splitter.create_documents([text])
            logger.info(f"Split text into {len(docs)} chunks")
            
            if prompt_type == "summarize":
                return self._process_summary(docs)
            elif prompt_type == "translate":
                return self._process_translation(docs, target_language)
            elif prompt_type == "extract_info":
                return self._process_extraction(docs)
            elif prompt_type == "qa":
                return self._process_qa(docs)
            else:
                return "不支持的处理类型"
        except Exception as e:
            logger.error(f"Error processing text: {str(e)}")
            raise

    def _process_summary(self, docs):
        """处理总结请求"""
        prompt = ChatPromptTemplate.from_template("""
            请总结以下文本，要求：
            1. 保持重要信息完整
            2. 条理清晰，结构化呈现
            3. 突出关键点
            4. 保持客观准确

            文本内容：
            {text}
        """)
        
        chain = (
            {"text": RunnablePassthrough()} 
            | prompt 
            | self.chat_model 
            | StrOutputParser()
        )
        
        # 处理每个文档块
        summaries = []
        for doc in docs:
            summary = chain.invoke(doc.page_content)
            summaries.append(summary)
        
        # 如果有多个总结，再次合并总结
        if len(summaries) > 1:
            final_prompt = ChatPromptTemplate.from_template("""
                请将以下多个段落的总结合并成一个完整的总结，要求：
                1. 保持逻辑连贯
                2. 去除重复内容
                3. 突出重要信息
                4. 结构清晰

                段落总结：
                {text}
            """)
            final_chain = (
                {"text": RunnablePassthrough()} 
                | final_prompt 
                | self.chat_model 
                | StrOutputParser()
            )
            return final_chain.invoke("\n\n".join(summaries))
        
        return summaries[0]

    def _process_translation(self, docs, target_language):
        """处理翻译请求"""
        prompt = ChatPromptTemplate.from_template("""
            请将以下文本翻译成{language}，要求：
            1. 保持专业术语的准确性
            2. 符合目标语言的表达习惯
            3. 保持原文的语气和风格
            4. 对专有名词保持一致的翻译

            文本内容：
            {text}
        """)
        
        chain = (
            {"text": RunnablePassthrough(), "language": lambda _: target_language}
            | prompt 
            | self.chat_model 
            | StrOutputParser()
        )
        
        results = []
        for doc in docs:
            result = chain.invoke(doc.page_content)
            results.append(result)
        return "\n".join(results)

    def _process_extraction(self, docs):
        """处理信息提取请求"""
        prompt = ChatPromptTemplate.from_template("""
            请从以下文本中提取关键信息，以结构化的方式呈现：

            要提取的内容包括：
            1. 主要观点和论述
            2. 重要数据和统计
            3. 时间、地点、人物
            4. 关键结论和建议
            5. 特殊术语和概念

            请使用 Markdown 格式输出，善用标题、列表等格式化元素。

            文本内容：
            {text}
        """)
        
        chain = (
            {"text": RunnablePassthrough()} 
            | prompt 
            | self.chat_model 
            | StrOutputParser()
        )
        
        results = []
        for doc in docs:
            result = chain.invoke(doc.page_content)
            results.append(result)
        return "\n".join(results)

    def _process_qa(self, docs):
        """处理问答生成请求"""
        prompt = ChatPromptTemplate.from_template("""
            请基于以下文本生成一个问答对列表，要求：
            1. 问题要有深度，覆盖文本的重要内容
            2. 答案要准确完整
            3. 问答对数量在5-10个之间
            4. 使用 Markdown 格式，每个问答对之间要有明确的分隔

            建议的问题类型：
            - 概念解释类
            - 原因分析类
            - 比较对比类
            - 应用实践类
            - 评价判断类

            文本内容：
            {text}
        """)
        
        chain = (
            {"text": RunnablePassthrough()} 
            | prompt 
            | self.chat_model 
            | StrOutputParser()
        )
        
        results = []
        for doc in docs:
            result = chain.invoke(doc.page_content)
            results.append(result)
        return "\n".join(results) 